import torch
import torch.nn as nn
from distribution import NormalSample


class BayesianBase(nn.Module):
    """贝叶斯基类
    Attributes:
        node: 随机节点
    """
    def __init__(self):
        super().__init__()
        self.node = {}

    def normal(self,
               name,
               mean,
               std,
               observation=None):
        """正态分布抽样函数
        Args:
            name ([str]): [节点名称]
            mean ([Tensor]): [正态分布均值参数, 也可以是 nn.Parameters(参数可以训练)]
            std ([Tensor]): [正态分布标准差参数, 也可以是 nn.Parameters(参数可以训练)]
            observation ([Tensor]): [分布的一组观测]. Defaults to None.
        """
        self.node[name] = NormalSample(mean, std, observation)

    @property
    def log_prob(self):
        """计算所有节点的条件概率密度
        Returns:
            [Tensor]: [description]
        """
        _log_prob = 0
        for key in self.node:
            _log_prob += self.node[key].log_prob()
        return _log_prob

    @property
    def get_samples(self, name: str):
        """[summary]

        Returns:
            [type]: [description]
        """
        self.node[name].sample()
        return self.node[name].get_samples
